#coding=utf-8
# 本文件必须依附 聚宽研究环境
import numpy as np
from datetime import datetime, date, timedelta
from dateutil.relativedelta import relativedelta
from jqdata import finance
import matplotlib.pyplot as plt
import pandas as pd

# data_df: 包含计算指标及比较的大盘指标两列的df
# min_v, max_v: 纵轴坐标范围，根据指数而定
def compareWithIndex(data_df:pd.DataFrame, min_v, max_v):
    plt.figure(figsize=(10,8))
    plt.plot(data_df)
    plt.ylim(min_v,max_v) # 纵轴坐标范围，根据指数而定.
# 蓝色是拥挤率*100， 橙色是指数


def max_DD(stock, end_date, count):
    stock_price = get_price(stock, end_date = end_date, count = count)['close']
    max_DD_list = np.zeros(count)
    for i in range(0, count-1):
        max_DD_list[i] = max(stock_price[i] - (stock_price[i:]))/stock_price[i]
    max_DD = max(max_DD_list)
    return max_DD


# MAD法去极值函数
def MAD(df, std_num=3):
    for i in df.columns:
        mad = np.median(np.absolute(df[i] - np.median(df[i])))
        ma = np.median(df[i])
        for x in df[i].index:
            if df[i][x] > (ma + std_num * 1.4826 * mad):
                df[i][x] = (ma + std_num * 1.4826 * mad)
            elif df[i][x] < (ma - std_num * 1.4826 * mad):
                df[i][x] = (ma - std_num * 1.4826 * mad)
    return df


# 因子标准化函数
def Standardize(df):
    for i in df.columns:
        df[i] = (df[i] - df[i].mean()) / df[i].std()
    return df


# 剔除ST股票
def removeSTstocks(stock_pool: list)->list:
    # 过滤ST等异常股
    stock_pool_temp = stock_pool.copy()
    for stock in stock_pool:
        info = finance.run_query(query(finance.STK_STATUS_CHANGE).filter(finance.STK_STATUS_CHANGE.code == stock))
        if 301003 in list(info.public_status_id.values):
            stock_pool_temp.remove(stock)

    return stock_pool_temp